import torch
import torchvision.datasets
from torch import nn
from torch.nn import Sigmoid

from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

#
# class JJw(nn.Module):
#     def __init__(self):
#         super(JJw,self).__init__()
#         self.maxpool1=MaxPool2d(kernel_size=3,ceil_mode=False)
#
#     def forward(self,input):
#         output=self.maxpool1(input)
#         return output
#
# dataset=torchvision.datasets.CIFAR10("./dataset",train=False,transform=torchvision.transforms.ToTensor(),download=True)
# dataloader=DataLoader(dataset,batch_size=64)
#
#
# # input=torch.tensor([[1,2,0,3,1]
# #                     ,[0,1,2,3,1]
# #                     ,[1,2,1,0,0]
# #                     ,[5,2,3,1,1]
# #                     ,[2,1,0,1,1]],dtype=torch.float32)
# # input=torch.reshape(input,(-1,1,5,5))
# jjw=JJw()
# step=0
# writer=SummaryWriter("logs")
#
# for data in dataloader:
#     imgs,targets=data
#     writer.add_images("input",imgs,step)
#     output=jjw(imgs)
#     writer.add_images("output",output,step)
#     step += 1
#
# writer.close()

class JJw(nn.Module):
    def __init__(self):
        super(JJw,self).__init__()
        # self.relu1=ReLU(inplace=False)
        self.sigmoid1=Sigmoid()

    def forward(self,input):
        output=self.sigmoid1(input)
        return output

dataset=torchvision.datasets.CIFAR10("./dataset",train=False,transform=torchvision.transforms.ToTensor(),download=True)
dataloader=DataLoader(dataset,batch_size=64)
jjw=JJw()
step=0
writer=SummaryWriter("logs")
for data in dataloader:
    imgs,targets=data
    writer.add_images("input",imgs,step)
    output=jjw(imgs)
    writer.add_images("output",output,step)
    step+=1
writer.close()
